#!/bin/bash
# AWS部署脚本 - 创建EC2实例并部署Turritopsis节点

set -e

# 配置参数
AWS_REGION="us-west-2"
INSTANCE_TYPE="t3.medium"
KEY_NAME="turritopsis-key"
SECURITY_GROUP="turritopsis-sg"
AMI_ID="ami-0c02fb55956c7d316"  # Amazon Linux 2 AMI
NODES_COUNT=9
PROJECT_DIR="/home/ec2-user/turritopsis"

# 颜色输出
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
NC='\033[0m' # No Color

log_info() {
    echo -e "${GREEN}[INFO]${NC} $1"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查AWS CLI是否安装
check_aws_cli() {
    if ! command -v aws &> /dev/null; then
        log_error "AWS CLI not found. Please install AWS CLI first."
        exit 1
    fi
    
    # 检查AWS配置
    if ! aws sts get-caller-identity &> /dev/null; then
        log_error "AWS credentials not configured. Please run 'aws configure' first."
        exit 1
    fi
    
    log_info "AWS CLI configured successfully"
}

# 创建安全组
create_security_group() {
    log_info "Creating security group..."
    
    # 检查安全组是否已存在
    if aws ec2 describe-security-groups --group-names "$SECURITY_GROUP" --region "$AWS_REGION" &> /dev/null; then
        log_warn "Security group $SECURITY_GROUP already exists"
        return
    fi
    
    # 创建安全组
    aws ec2 create-security-group \
        --group-name "$SECURITY_GROUP" \
        --description "Security group for Turritopsis nodes" \
        --region "$AWS_REGION" \
        --output text --query 'GroupId'
    
    # 添加SSH规则
    aws ec2 authorize-security-group-ingress \
        --group-name "$SECURITY_GROUP" \
        --protocol tcp \
        --port 22 \
        --cidr 0.0.0.0/0 \
        --region "$AWS_REGION"
    
    # 添加Turritopsis端口规则
    aws ec2 authorize-security-group-ingress \
        --group-name "$SECURITY_GROUP" \
        --protocol tcp \
        --port 19001-19020 \
        --cidr 0.0.0.0/0 \
        --region "$AWS_REGION"
    
    log_info "Security group created successfully"
}

# 创建EC2实例
create_instances() {
    log_info "Creating $NODES_COUNT EC2 instances..."
    
    # 创建实例
    INSTANCE_IDS=$(aws ec2 run-instances \
        --image-id "$AMI_ID" \
        --count "$NODES_COUNT" \
        --instance-type "$INSTANCE_TYPE" \
        --key-name "$KEY_NAME" \
        --security-groups "$SECURITY_GROUP" \
        --region "$AWS_REGION" \
        --tag-specifications "ResourceType=instance,Tags=[{Key=Name,Value=turritopsis-node},{Key=Project,Value=Turritopsis}]" \
        --output text --query 'Instances[].InstanceId')
    
    log_info "Created instances: $INSTANCE_IDS"
    
    # 等待实例运行
    log_info "Waiting for instances to be running..."
    aws ec2 wait instance-running --instance-ids $INSTANCE_IDS --region "$AWS_REGION"
    
    # 获取实例IP地址
    aws ec2 describe-instances \
        --instance-ids $INSTANCE_IDS \
        --region "$AWS_REGION" \
        --query 'Reservations[].Instances[].[InstanceId,PublicIpAddress,PrivateIpAddress]' \
        --output table
    
    # 保存实例信息到文件
    aws ec2 describe-instances \
        --instance-ids $INSTANCE_IDS \
        --region "$AWS_REGION" \
        --query 'Reservations[].Instances[].[InstanceId,PublicIpAddress,PrivateIpAddress]' \
        --output text > aws_instances.txt
    
    log_info "Instance information saved to aws_instances.txt"
}

# 生成hosts文件
generate_hosts_file() {
    log_info "Generating hosts file..."
    
    if [ ! -f "aws_instances.txt" ]; then
        log_error "aws_instances.txt not found. Please run create_instances first."
        exit 1
    fi
    
    # 生成hosts文件
    cat > aws_hosts.txt << EOF
# Turritopsis AWS Nodes
# Generated on $(date)

EOF
    
    node_id=0
    while read -r instance_id public_ip private_ip; do
        if [ "$public_ip" != "None" ] && [ "$private_ip" != "None" ]; then
            echo "node-$node_id $public_ip $private_ip" >> aws_hosts.txt
            ((node_id++))
        fi
    done < aws_instances.txt
    
    log_info "Hosts file generated: aws_hosts.txt"
}

# 部署代码到所有节点
deploy_to_nodes() {
    log_info "Deploying code to all nodes..."
    
    if [ ! -f "aws_hosts.txt" ]; then
        log_error "aws_hosts.txt not found. Please run generate_hosts_file first."
        exit 1
    fi
    
    # 创建部署脚本
    cat > deploy_node.sh << 'EOF'
#!/bin/bash
set -e

# 更新系统
sudo yum update -y

# 安装Go
if ! command -v go &> /dev/null; then
    curl -O https://golang.org/dl/go1.21.0.linux-amd64.tar.gz
    sudo tar -C /usr/local -xzf go1.21.0.linux-amd64.tar.gz
    echo 'export PATH=$PATH:/usr/local/go/bin' >> ~/.bashrc
    echo 'export GOPATH=$HOME/go' >> ~/.bashrc
    source ~/.bashrc
fi

# 创建项目目录
mkdir -p $PROJECT_DIR
cd $PROJECT_DIR

# 克隆代码（这里假设代码已经在本地）
# 实际部署时，需要将代码上传到S3或使用git clone

# 构建项目
go mod tidy
go build ./...

# 创建运行目录
mkdir -p run

echo "Deployment completed on $(hostname)"
EOF
    
    # 分发部署脚本到所有节点
    while read -r node_name public_ip private_ip; do
        if [[ "$node_name" =~ ^node- ]]; then
            log_info "Deploying to $node_name ($public_ip)..."
            
            # 复制部署脚本
            scp -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no deploy_node.sh ec2-user@$public_ip:/tmp/
            
            # 执行部署
            ssh -i ~/.ssh/$KEY_NAME.pem -o StrictHostKeyChecking=no ec2-user@$public_ip \
                "chmod +x /tmp/deploy_node.sh && PROJECT_DIR=$PROJECT_DIR /tmp/deploy_node.sh"
        fi
    done < aws_hosts.txt
    
    log_info "Deployment completed on all nodes"
}

# 主函数
main() {
    log_info "Starting AWS deployment for Turritopsis..."
    
    check_aws_cli
    create_security_group
    create_instances
    generate_hosts_file
    deploy_to_nodes
    
    log_info "AWS deployment completed successfully!"
    log_info "Next steps:"
    log_info "1. Review aws_hosts.txt for node information"
    log_info "2. Run ./aws_start_turritopsis.sh to start the protocol"
    log_info "3. Run ./aws_collect_metrics.sh to collect results"
}

# 脚本参数处理
case "${1:-}" in
    "create-sg")
        check_aws_cli
        create_security_group
        ;;
    "create-instances")
        check_aws_cli
        create_instances
        ;;
    "generate-hosts")
        generate_hosts_file
        ;;
    "deploy")
        deploy_to_nodes
        ;;
    "cleanup")
        log_info "Cleaning up AWS resources..."
        if [ -f "aws_instances.txt" ]; then
            INSTANCE_IDS=$(awk '{print $1}' aws_instances.txt | tr '\n' ' ')
            aws ec2 terminate-instances --instance-ids $INSTANCE_IDS --region "$AWS_REGION"
            log_info "Instances terminated"
        fi
        ;;
    *)
        main
        ;;
esac
